#include "FileSender.h"
#include "MD5.h"
#include <sys/socket.h>
#include <arpa/inet.h>
#include <sys/stat.h>
#include <unistd.h>
#include <cstring>
#include <cstdio>
#include <cstdlib>
#include <netinet/in.h>
#include <sys/time.h>
#include <mutex>
#include <condition_variable>
#include <atomic>

#define CHUNK_SIZE 4096

// 静态成员初始化
std::mutex FileSender::mtx_;
std::condition_variable FileSender::cv_;
std::atomic<bool> FileSender::start_signal_(false);
std::atomic<bool> FileSender::abort_signal_(false);

FileSender::FileSender(const std::string& file_path, const std::string& dest_ip, int dest_port, int local_port)
    : file_path_(file_path), dest_ip_(dest_ip), dest_port_(dest_port), local_port_(local_port), buffer_(nullptr), filesize_(0) {
    loadFile();
}

void FileSender::loadFile() {
    FILE* fp = fopen(file_path_.c_str(), "rb");
    if (!fp) {
        perror("Failed to open file");
        exit(1);
    }
    struct stat st;
    if (stat(file_path_.c_str(), &st) != 0) {
        perror("stat error");
        fclose(fp);
        exit(1);
    }
    filesize_ = st.st_size;
    buffer_ = (unsigned char*)malloc(filesize_);
    if (!buffer_) {
        perror("malloc error");
        fclose(fp);
        exit(1);
    }
    fread(buffer_, 1, filesize_, fp);
    fclose(fp);
}

void FileSender::sendUDP() {
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket error");
        exit(1);
    }
    int bufsize = 8 * 1024 * 1024;
    setsockopt(sock, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
    setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));
    if (local_port_ > 0) {
        struct sockaddr_in local_addr;
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port_);
        if (bind(sock, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind local port error");
            close(sock);
            exit(1);
        }
    }
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port_);
    inet_pton(AF_INET, dest_ip_.c_str(), &dest_addr.sin_addr);
    size_t total_sent = 0;
    unsigned int seq = 0;
    while (total_sent < filesize_) {
        size_t chunk = (filesize_ - total_sent > CHUNK_SIZE) ? CHUNK_SIZE : (filesize_ - total_sent);
        unsigned char sendbuf[CHUNK_SIZE + 4];
        memcpy(sendbuf, &seq, 4);
        memcpy(sendbuf + 4, buffer_ + total_sent, chunk);
        ssize_t sent = sendto(sock, sendbuf, chunk + 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
        if (sent < 0) {
            perror("sendto error");
            break;
        }
        total_sent += chunk;
        seq++;
        usleep(1);
    }
    unsigned int end_seq = 0xFFFFFFFF;
    sendto(sock, &end_seq, 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
    close(sock);
}

// 新增：多线程多文件UDP发送，包头1字节文件编号+4字节序号
void FileSender::sendUDP(int file_id) {
    while (true) {
        {
            std::unique_lock<std::mutex> lk(mtx_);
            cv_.wait(lk, []{ return start_signal_ || abort_signal_; });
            if (abort_signal_) {
                // 释放内存、重置变量
                if (buffer_) { free(buffer_); buffer_ = nullptr; }
                filesize_ = 0;
                printf("[FileSender][file_id=%d] ABORT received before start, waiting...\n", file_id);
                continue; // 被ABORT打断，回到等待START
            }
        }
    printf("[FileSender][file_id=%d] Start sending file: %s, size: %zu bytes, to %s:%d\n", file_id, file_path_.c_str(), filesize_, dest_ip_.c_str(), dest_port_);
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("socket error");
        printf("[FileSender][file_id=%d] Socket creation failed!\n", file_id);
        return;
    }
    int bufsize = 8 * 1024 * 1024;
    setsockopt(sock, SOL_SOCKET, SO_SNDBUF, &bufsize, sizeof(bufsize));
    setsockopt(sock, SOL_SOCKET, SO_RCVBUF, &bufsize, sizeof(bufsize));
    if (local_port_ > 0) {
        struct sockaddr_in local_addr;
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port_);
        if (bind(sock, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind local port error");
            printf("[FileSender][file_id=%d] Bind local port %d failed!\n", file_id, local_port_);
            close(sock);
            return;
        }
    }
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port_);
    inet_pton(AF_INET, dest_ip_.c_str(), &dest_addr.sin_addr);
    size_t total_sent = 0;
    unsigned int seq = 0;
    bool aborted = false;
    while (total_sent < filesize_) {
        if (abort_signal_) {
            // 释放内存、重置变量
            if (buffer_) { free(buffer_); buffer_ = nullptr; }
            filesize_ = 0;
            aborted = true;
            printf("[FileSender][file_id=%d] ABORT received during sending, interrupted at seq=%u, sent=%zu bytes.\n", file_id, seq, total_sent);
            break;
        }
        size_t chunk = (filesize_ - total_sent > CHUNK_SIZE) ? CHUNK_SIZE : (filesize_ - total_sent);
        unsigned char sendbuf[CHUNK_SIZE + 1 + 4];
        sendbuf[0] = static_cast<unsigned char>(file_id);
        memcpy(sendbuf + 1, &seq, 4);
        memcpy(sendbuf + 1 + 4, buffer_ + total_sent, chunk);
        ssize_t sent = sendto(sock, sendbuf, chunk + 1 + 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
        if (sent < 0) {
            perror("sendto error");
            printf("[FileSender][file_id=%d] sendto failed at seq=%u, sent=%zu bytes.\n", file_id, seq, total_sent);
            break;
        }
        total_sent += chunk;
        seq++;
        usleep(1);
    }
    // 发送完成或被ABORT后，发送结束包
    unsigned char endbuf[1 + 4];
    endbuf[0] = static_cast<unsigned char>(file_id);
    unsigned int end_seq = 0xFFFFFFFF;
    memcpy(endbuf + 1, &end_seq, 4);
    sendto(sock, endbuf, 1 + 4, 0, (struct sockaddr*)&dest_addr, sizeof(dest_addr));
    close(sock);
    if (!aborted) {
        printf("[FileSender][file_id=%d] File send complete! Total sent: %zu bytes, seq=%u\n", file_id, total_sent, seq);
    }
    {
        std::lock_guard<std::mutex> lk(mtx_);
        start_signal_ = false;
    }
    }
}

void FileSender::sendTCP() {
    int sock = socket(AF_INET, SOCK_STREAM, 0);
    if (sock < 0) {
        perror("socket error");
        exit(1);
    }
    if (local_port_ > 0) {
        struct sockaddr_in local_addr;
        memset(&local_addr, 0, sizeof(local_addr));
        local_addr.sin_family = AF_INET;
        local_addr.sin_addr.s_addr = htonl(INADDR_ANY);
        local_addr.sin_port = htons(local_port_);
        if (bind(sock, (struct sockaddr*)&local_addr, sizeof(local_addr)) < 0) {
            perror("bind local port error");
            close(sock);
            exit(1);
        }
    }
    struct sockaddr_in dest_addr;
    memset(&dest_addr, 0, sizeof(dest_addr));
    dest_addr.sin_family = AF_INET;
    dest_addr.sin_port = htons(dest_port_);
    inet_pton(AF_INET, dest_ip_.c_str(), &dest_addr.sin_addr);
    if (connect(sock, (struct sockaddr*)&dest_addr, sizeof(dest_addr)) < 0) {
        perror("connect error");
        close(sock);
        exit(1);
    }
    uint64_t fsize = (uint64_t)filesize_;
    send(sock, &fsize, sizeof(fsize), 0);
    size_t total_sent = 0;
    while (total_sent < filesize_) {
        size_t chunk = (filesize_ - total_sent > CHUNK_SIZE) ? CHUNK_SIZE : (filesize_ - total_sent);
        ssize_t sent = send(sock, buffer_ + total_sent, chunk, 0);
        if (sent < 0) {
            perror("send error");
            break;
        }
        total_sent += sent;
    }
    unsigned char md5[16];
    MD5 ctx;
    ctx.update(buffer_, filesize_);
    ctx.final(md5);
    send(sock, md5, 16, 0);
    close(sock);
}

// 新增：监听127.0.0.1:9000，收到START指令唤醒所有线程
void FileSender::waitForCommand() {
    int sock = socket(AF_INET, SOCK_DGRAM, 0);
    if (sock < 0) {
        perror("command socket error");
        exit(1);
    }
    struct sockaddr_in addr;
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = inet_addr("127.0.0.1");
    addr.sin_port = htons(9000);
    if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) < 0) {
        perror("command bind error");
        close(sock);
        exit(1);
    }
    char buf[128];
    while (true) {
        ssize_t n = recvfrom(sock, buf, sizeof(buf)-1, 0, NULL, NULL);
        if (n > 0) {
            buf[n] = '\0';
            std::string cmd(buf);
            if (cmd == "START") {
                {
                    std::lock_guard<std::mutex> lk(mtx_);
                    start_signal_ = true;
                    abort_signal_ = false;
                }
                cv_.notify_all();
            } else if (cmd == "ABORT") {
                abort_signal_ = true;
            }
        }
    }
    close(sock);
} 